package org.xq.softcup.mq.remoting.invoker.router.impl;

import org.xq.softcup.mq.remoting.invoker.router.XQRpcLoadBalance;

import java.io.UnsupportedEncodingException;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.SortedMap;
import java.util.TreeMap;
import java.util.TreeSet;

public class XQRpcLoadBalanceConsistentHashStrategy extends XQRpcLoadBalance {
    private int VIRTUAL_NODE_NUM = 5;

    private long hash(String key)
    {
        MessageDigest md5;
        try
        {
            md5 = MessageDigest.getInstance("MD5");
        }
        catch (NoSuchAlgorithmException e)
        {
            throw new RuntimeException("MD5 not supported", e);
        }
        md5.reset();
        byte[] keyBytes = null;
        try
        {
            keyBytes = key.getBytes("UTF-8");
        }
        catch (UnsupportedEncodingException e)
        {
            throw new RuntimeException("Unknown string :" + key, e);
        }
        md5.update(keyBytes);
        byte[] digest = md5.digest();
        long hashCode = (digest[3] & 0xFF) << 24 | (digest[2] & 0xFF) << 16 | (digest[1] & 0xFF) << 8 | digest[0] & 0xFF;
        long truncateHashCode = hashCode & 0xFFFFFFFF;
        return truncateHashCode;
    }

    public String doRoute(String serviceKey, TreeSet<String> addressSet)
    {
        TreeMap<Long, String> addressRing = new TreeMap();
        for (String address : addressSet) {
            for (int i = 0; i < this.VIRTUAL_NODE_NUM; i++)
            {
                long addressHash = hash("SHARD-" + address + "-NODE-" + i);
                addressRing.put(Long.valueOf(addressHash), address);
            }
        }
        long jobHash = hash(serviceKey);
        SortedMap<Long, String> lastRing = addressRing.tailMap(Long.valueOf(jobHash));
        if (!lastRing.isEmpty()) {
            return (String)lastRing.get(lastRing.firstKey());
        }
        return (String)addressRing.firstEntry().getValue();
    }

    public String route(String serviceKey, TreeSet<String> addressSet)
    {
        String finalAddress = doRoute(serviceKey, addressSet);
        return finalAddress;
    }
}
